import hashlib
import os
import queue
import subprocess
import sys
import threading


class xrayFileScan:
    def __init__(self):
        self.url_q = queue.Queue()
        self.new_file_path = ""
        self.lock = threading.Lock()

    def xrayScan(self, url):
        if url[0:4] not in "http":
            url = "http://{}".format(url)
        outputfilename = hashlib.md5(url.encode("utf-8"))
        scanCommand = "./xray webscan --basic-crawler {} --html-output {}/{}.html".format(url, self.new_file_path,
                                                                                            outputfilename.hexdigest())
        print(scanCommand)
        subprocess.run(scanCommand, shell=True)

    def banner(self):
        banner_info = '''
                 _____ _ _     _____             
 _ _ ___ ___ _ _|   __|_| |___|   __|___ ___ ___ 
|_'_|  _| .'| | |   __| | | -_|__   |  _| .'|   |
|_,_|_| |__,|_  |__|  |_|_|___|_____|___|__,|_|_|
            |___|           
                             ________________________                         
                            | author: Spike.         |
                            |  email: lendme@163.com |
                            |________________________| 
    '''
        print(banner_info)

    def main(self):
        self.banner()
        url_files = ""
        threading_count = ""
        if len(sys.argv) == 2:
            url_files = sys.argv[1]
            if "-h" in url_files:
                print("""
-------------------------------------------------------------
usage:  
    python3 xrayFileScan.py url文件 线程数量(默认开启3个线程!)
    python3 xrayFileScan.py target.txt 5

                """)
                sys.exit()
            # 如果不设置进程参数，则默认开启3个进程
            threading_count = 3
        if len(sys.argv) == 3:
            url_files = sys.argv[1]
            threading_count = int(sys.argv[2])
        url_files_name = url_files[:-4]
        current_path = os.getcwd()
        os.mkdir(url_files_name)
        self.new_file_path = "{}/{}".format(current_path, url_files_name)
        with open(url_files, "r", encoding="utf-8") as url_files:
            for url in url_files:
                url = url.strip()
                self.url_q.put(url)
            self.get_value(threading_count)

    def get_value(self, threading_count):
        threads = []
        for i in range(1, threading_count + 1):
            t = threading.Thread(target=self.worker)
            t.start()
            threads.append(t)
        for t in threads:
            t.join()

    def worker(self):
        while True:
            with self.lock:
                if not self.url_q.empty():
                    url = self.url_q.get()
                else:
                    break
            self.xrayScan(url)


if __name__ == '__main__':
    run = xrayFileScan()
    run.main()
